'''
Author: zhuyuejiang
Date: 2021-04-25 00:57:30
LastEditTime: 2021-05-19 02:15:35
Description:
行为分析算法处理器,用于从指定算法包中导入行为分析算法,并调用指定接口,运行行为分析算法,
最后讲算法运行信息以及运行结果反馈至API SERVER
'''


import os
import time
import cv2
import yaml
import sys
import unicodedata
from yacs.config import CfgNode
import os.path as osp
import importlib
import FrameWork.utils as action_utils
import math
import shutil


'''
@description: 
行为分析算法处理器类
@param {*}
@return {*}
'''
class ActionProcessor:
    def __init__(self):
        self._workdir = os.getcwd()        #用于记录主程序运行目录
        self._module_name = ""
        self._error_string = "OK"

    '''
    @description: 任务初始化，首先是从指定配置文件读取配置参数，如果配置
                  文件不存在，则默认从算法库目录下寻找名为‘action_cfg.yaml’
                  的配置文件，寻找失败，则采用默认配置参数；然后设置输出路径；
                  接着是从‘library_path'下导入算法模块；导入完成后，启动算法
                  运行流程
    @param {*} self
    @param {*} library_path
    @param {*} output_path
    @param {*} data_path
    @param {*} config_file
    @return {*}
    '''
    def init_job(self, library_path, output_path, data_path, config_file):
        cfg_dict = {}
        cfg_file = library_path + "/action_cfg.yaml"
        if len(config_file) > 0:
            cfg_file = config_file
        print("[Info:] config file:", cfg_file)
        with open(cfg_file, 'r') as fo:
            cfg_dict.update(yaml.load(fo.read(), yaml.FullLoader))
        cfg_yaml = CfgNode(cfg_dict)
        self._module_name = cfg_yaml.GeneralOptions.ModelName
        if len(data_path) > 0:
            cfg_yaml.GeneralOptions.PATH_TO_DATA_DIR = data_path
        
        # import lib
        if sys.modules.get('actionRecognition') is not None:
            del sys.modules['actionRecognition']
        if sys.modules.get('actionRecognition.actionRecognition') is not None:
            del sys.modules['actionRecognition.actionRecognition']
        importlib.invalidate_caches()
        action_lib = importlib.import_module('actionRecognition.actionRecognition')
        while len(sys.argv) > 1 :
            sys.argv.pop()
        ar = action_lib.ZJLAB_ACTIONRECOGNITION(cfg_yaml)

        action_utils.mkdir_if_missing(output_path) #创建输出目录

        return ar

    '''
    @description: 用于判断输入字符串是否为数字
    @param {*}
    @return {*}
    '''
    @staticmethod
    def is_number(str):
        try:
            if 'NaN' == str or str is None:
                return False
            float(str)
            return True
        except ValueError:
            pass

        try:
            unicodedata.numeric(str)
        except(TypeError, ValueError):
            pass

        return False

    '''
    @description:  保存一帧的处理结果
    @param {*} self
    @param {*} frame
    @param {*} pred
    @param {*} frame_index
    @param {*} output_path
    @return {*}
    '''
    def save_frame(self, frame, pred, frame_index, output_path):
        y_offset = 50
        clr = (0, 0, 255)
        title = "Action:"
        label = "unkown"
        if ActionProcessor.is_number(pred):
            title = "CrowdCount:"
            label = '{:.0f}'.format(pred)
        elif not pred is None:
            label = '{}'.format(pred)
            
        
        cv2.putText(
            frame,
            title,
            (10, y_offset),
            fontFace = cv2.FONT_HERSHEY_SIMPLEX,
            fontScale = 0.65,
            color= clr,
            thickness=2,
        )

        y_offset += 30
        print("[Info:] label:", label)
        cv2.putText(
            frame,
            label,
            (20, y_offset),
            fontFace = cv2.FONT_HERSHEY_SIMPLEX,
            fontScale = 0.65,
            color = clr,
            thickness= 2,
        )
        cv2.imwrite(osp.join(output_path, 'frame' + str(frame_index) + '.jpg'), frame)


    '''
    @description: 在数据集上运行行为分析算法
    @param {*} self
    @param {*} library_path
    @param {*} input_path
    @param {*} output_path
    @param {*} config_file
    @return {*}
    '''
    def run_eval(self, library_path, data_path, result_path, config_file):
        print("action task starting-----")
        #切换到算法包所在路径
        sys.path.append(library_path)
        os.chdir(library_path)
        
        #算法初始化
        ar = self.init_job(library_path, result_path, data_path, config_file)

        #内置的crowdcount算法运行模式与其余行为分析算法不同，不可在数据集上运行
        if self._module_name.startswith('crowdcount'):
            self._error_string = "CrowdCount can not run in dataset model..."
            print("[Err:]", self._error_string)
        
        #运行算法
        video_names, pred_labels, top1_acc, top5_acc = ar.run(None)
        print("[Info:] Action recognition in %s finished...." % data_path)
        print("PlatformProcessNum:50")
        #每段视频中截取一张图,打上预测结果标签,保存到result_path中
        label_index = 0
        for video in video_names:
            capture = cv2.VideoCapture(video)
            total_frame = capture.get(cv2.CAP_PROP_FRAME_COUNT)
            target_frame = math.floor(total_frame / 2)
            capture.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
            ret, frame = capture.read()
            if ret :
                self.save_frame(frame, pred_labels[label_index], label_index + 1, result_path)
            label_index += 1
            print("PlatformProcessNum:", int(50 + label_index * 50 / len(video_names)))
        sys.path.remove(library_path)
        print("PlatformSummary:", '{"top1_acc":%s, "top5_acc":%s}' % (top1_acc, top5_acc))


if __name__ == '__main__':
    print("[Info:] actionrunner recieved task, action recognition process starting............................................")
    print("[Info:] --------------program name:", sys.argv[0])
    print("[Info:] --------------library path:", sys.argv[1])
    print("[Info:] --------------model file:", sys.argv[2])
    print("[Info:] --------------config file:", sys.argv[3])
    print("[Info:] --------------dataset path:", sys.argv[4])
    print("[Info:] --------------result root dir:", sys.argv[5])
    print("[Info:] --------------data type:", sys.argv[6])
    
    if 0 == int(sys.argv[6]):
        print("[Err:] sorry, video is unsupported by action recognition algorithms")
    else:
        store_path = "data/result/ar"
        if len(sys.argv[5]) > 0:
            store_path = sys.argv[5]
        actionProc = ActionProcessor()
        actionProc.run_eval(library_path=str(sys.argv[1]), data_path=str(sys.argv[4]), result_path=str(store_path), config_file=str(sys.argv[3]))
    